# -*- coding: utf-8 -*-
# 本程序用于视频中车辆行人等多目标检测跟踪
# @Time    : 2022/7/13 10:29
# @Author  : ShenYue
# @Email   : 19221227@bjtu.edu.cn
# @Software: PyCharm

import warnings
from collections import deque
import numpy as np
import imutils
import time
import cv2
import os
from tqdm import tqdm
from sort import Sort
from Detector_Model import Detector

warnings.filterwarnings('ignore')

class Car:
    def __init__(self,c_id,c_x,c_y,direction,c_count,c_start_time,c_over_time,c_v):
        self.c_id = c_id    # 车辆编号
        self.c_x=c_x    # 中心点坐标
        self.c_y=c_y
        self.direction = direction  # 上下行方向
        self.c_count = c_count  # 车辆计数？
        self.c_start_time = c_start_time    # 通过检测线时间长短
        self.c_over_time = c_over_time
        self.c_v = c_v
    def updateCoords(self,x,y):
        self.c_x= x
        self.c_y= y

def Feature_calculate(frame,thiscap,dets,cars,line_height,bound,pid,offset,frame_width):
    (line_height1,line_height2)=line_height
    (left_bound1,left_bound2,right_bound1,right_bound2)=bound
    add_count_up = 0
    add_count_down = 0  # 车辆计数
    # add_occup_time = 0  # 线占用时间
    v=0

    # 定义车辆参数计算规则
    cv2.line(frame, (left_bound1, line_height1), (right_bound1, line_height1), (255, 255, 255), 2)
    cv2.line(frame, (left_bound2, line_height2), (right_bound2, line_height2), (255, 255, 255), 2)
    cv2.line(frame, (0, line_height1), (frame_width, line_height1), (255, 255, 0), 1)
    cv2.line(frame, (0, line_height2), (frame_width, line_height2), (255, 255, 0), 1)

    null = True #如果这一帧没有检测到任何车，则将cars置空，减小误差,设置初始变量null设为True
    for item in dets:
        (x, y, w, h) = int(item[0]), int(item[1]), int(item[2]-item[0]),int(item[3]-item[1])#dets:[x, y, x + w, y + h]
        # 计算大轮廓的中心点,也就是矩形框的中心点，并把中心点都放到cars里面
        cx = int(x + w / 2)
        cy = int(y + h / 2)

        if line_height2-30<int(y+(h/2))<line_height1+30 and left_bound1<int(x+(w/2))<right_bound2:
            null = False #有车时不会清空
            new = True  #是否创建车辆
            for i in cars:
                if  abs(cx-i.c_x) <100 and abs(cy - i.c_y)<100:#找到这辆车与上一帧中最近的车
                    # cv.putText(frame, 'cid:' + str(i.c_id), (x, y-10), cv.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 2)
                    # cv.putText(frame, "now_v:" + str('%.4f' % (i.c_v)) + 'm/s', (x, y-10),
                    #            cv.FONT_HERSHEY_COMPLEX, 1,
                    #            (255, 0, 0), 2)
                    new = False
                    i.updateCoords(cx,cy)#更新车辆位置信息
                    cv2.putText(frame, "vehicle-"+str(i.c_id) , (int(x + w / 2), int(y - 5)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)

                    # line_height1-offset<=
                    if i.c_y>=line_height1+offset and i.direction =='down' and i.c_count ==False:
                        i.c_over_time = thiscap
                        i.c_v = distance*3.6*fps/(i.c_over_time - i.c_start_time)
                        if i.c_v>v:
                            v = i.c_v
                        add_count_down+=1
                        i.c_count=True
                       # 向界面输出速度
                    # <=line_height2+offset
                    if line_height2-offset>=i.c_y and i.direction =='up' and (i.c_count == False):
                        i.c_over_time = thiscap
                        i.c_v = distance*3.6*fps/(i.c_over_time - i.c_start_time)
                        if i.c_v > v:
                            v = i.c_v
                        add_count_up+=1
                        i.c_count=True
                        # 向界面输出速度

                # if i.c_y>720 or i.c_y<380:
                #     cars.remove(i) #超过一定范围，删除对象
            if new == True : #符合一定条件，创建对象,and line_height1-offset<cy<line_height2+offset
                start_time = thiscap
                p = Car(pid, cx, cy, 'unknow', False, start_time ,0,0)
                if p.c_x <left_bound2:
                    p.direction = 'down'
                else:
                    p.direction = 'up'
                cars.append(p)
                pid += 1
            cv2.circle(frame, (cx, cy), 5, (0, 0, 255), -1)
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 3)

    if null == True:    #该帧没车，清空cars
        cars = []
    return add_count_up,add_count_down,v,cars

def Feafure_show(frame,count_up,count_down,max_v):
    cv2.putText(frame,"Up:"+str(count_up),(100, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255),5)
    cv2.putText(frame, "Down:" + str(count_down),(100, 100), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 5)
    cv2.putText(frame, "max_v:" + str('%.2f' %(max_v*3.6))+'km/h', (100, 130), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 5)
    # cv2.putText(frame, 'Occupancy Ratio:' + '%.2f' % (occup_time * 100 / thiscap) + '%', (100, 130),
    #             cv2.FONT_HERSHEY_SIMPLEX, 1,
    #             (255, 0, 255), 5)  # 第18章知识点，讲计数结果显示在视频上




if __name__ == '__main__':
##----------------交通参量定义------------------------------##
    carnum=0
    occup_time=0
    count_up = 0
    count_down = 0
    cars = []
    pid = 1
    max_v = 0
    distance = 6

    line_height=(570,470)
    # line_height1 = 570  # 红计数线位置
    # line_height2 = 470  # 红计数线位置
    offset = 7  # 计数线上下的偏移量
    bound=(0,560,660,1000)
    # left_bound1 = 0
    # right_bound1 = 560
    # left_bound2 = 660
    # right_bound2 = 1000

##----------------视频检测模型初始化------------------------------##
    model_path = "./yolo-obj/yolov4.weights"  # 模型文件路径
    detector_model = Detector(model_path=model_path)

    # 初始化用于标记框的颜色
    np.random.seed(42)
    COLORS = np.random.randint(0, 255, size=(200, 3), dtype="uint8")
    # dets, boxes, indexIDs, cls_IDs = detector_model.run(image)  # 执行检测跟踪
    # 载入数据集标签
    labelsPath = "./yolo-obj/coco.names"
    LABELS = open(labelsPath).read().strip().split("\n")

##----------------离线视频流读取------------------------------##
    video_path = "./video/highway.mp4"  # 要检测的视频路径
    vs = cv2.VideoCapture(video_path)
    fps = round(vs.get(cv2.CAP_PROP_FPS))
    (W, H) = (None, None)
    thiscap = 0  # 视频帧数
    ret, frame = vs.read()
    vw = frame.shape[1]
    vh = frame.shape[0]
    print("[INFO] 视频尺寸：{} * {}".format(vw, vh))
    fourcc = cv2.VideoWriter_fourcc(*'XVID')  # 该参数是MPEG-4编码类型，文件名后缀为.avi，可指定结果视频的大小
    output_video = cv2.VideoWriter(video_path.replace(".mp4", "-det.avi"), fourcc, 20.0, (vw, vh))  # 处理后的视频对象

##---------------------逐帧检测---------------------------##

    # 遍历视频帧进行检测，转换成进度条模式
    while(True):
        # 实时显示检测画面
        (grabbed, frame) = vs.read()
        if not grabbed:# 若grabbed为空，表示视频到达最后一帧，退出
            break
        if W is None or H is None:# 获取画面长宽
            (H, W) = frame.shape[:2]
        dets, boxes, indexIDs, cls_IDs,det_cars = detector_model.run(frame)# dets:[x, y, x + w, y + h, confidences[i], classIDs[i]]
        if len(boxes) > 0:
            i = int(0)
            for box in boxes:  # 遍历所有标记框
                (x, y) = (int(box[0]), int(box[1]))
                (w, h) = (int(box[2]), int(box[3]))

                # 在图像上标记目标框
                color = [int(c) for c in COLORS[indexIDs[i] % len(COLORS)]]
                cv2.rectangle(frame, (x, y), (w, h), color, 4)
                center = (int(((box[0]) + (box[2])) / 2), int(((box[1]) + (box[3])) / 2))
                # 显示某个对象标记框的中心
                thickness = 5
                cv2.circle(frame, center, 1, color, thickness)
                # 标记跟踪到的目标和数目
                text = "{}-{}".format(LABELS[int(cls_IDs[i])], indexIDs[i])
                cv2.putText(frame, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3)
                i += 1

##-------------------车辆计数----------------------------##
        add_count_up, add_count_down,add_occup_time,v,cars= Feature_calculate(frame,thiscap,det_cars,cars,line_height,bound,pid,offset,W)
        count_up=count_up+add_count_up
        count_down=count_down+add_count_down
        max_v=max(max_v,v)
        occup_time=occup_time+add_occup_time

        Feafure_show(frame,count_up,count_down,max_v)
        cv2.imshow('Stream', frame)#在GUI里显示一幅图像
        output_video.write(frame)  # 保存标记后的视频
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        #计算检测时间
        # print("FPS:{}".format(int(0.6/(end-start))))
        thiscap += 1

        # if thiscap >= total:  # 可设置检测的最大帧数提前退出
        #     print("[INFO] 运行结束...")
        #     output_video.release()
        #     vs.release()
        #     exit()
